#include <iostream>
#include <cmath>
#include <cuComplex.h>

// FFT kernel - 扩展为处理三组独立数据
__global__ void fft_kernel(cuComplex* data, int N) {
    int idx = threadIdx.x + blockIdx.x * blockDim.x;
    
    // 第一组计算 (0*N to 1*N-1)
    if (idx < N) {
        cuComplex temp = make_cuComplex(0.0f, 0.0f);
        for (int k = 0; k < N; k++) {
            float realPart = cosf(-2.0f * M_PI * idx * k / N);
            float imagPart = sinf(-2.0f * M_PI * idx * k / N);
            cuComplex twiddle = make_cuComplex(realPart, imagPart);
            temp = cuCaddf(temp, cuCmulf(data[k], twiddle));
        }
        data[idx] = temp;
    }
    // 第二组计算 (1*N to 2*N-1)
    else if (idx >= N && idx < 2 * N) {
        int m = idx % N;
        cuComplex temp = make_cuComplex(0.0f, 0.0f);
        for (int k = 0; k < N; k++) {
            float realPart = cosf(-2.0f * M_PI * m * k / N);
            float imagPart = sinf(-2.0f * M_PI * m * k / N);
            cuComplex twiddle = make_cuComplex(realPart, imagPart);
            temp = cuCaddf(temp, cuCmulf(data[N + k], twiddle));
        }
        data[idx] = temp;
    }
    // 第三组计算 (2*N to 3*N-1) - 新增的冗余模块
    else if (idx >= 2 * N && idx < 3 * N) {
        int m = idx % N;
        cuComplex temp = make_cuComplex(0.0f, 0.0f);
        for (int k = 0; k < N; k++) {
            float realPart = cosf(-2.0f * M_PI * m * k / N);
            float imagPart = sinf(-2.0f * M_PI * m * k / N);
            cuComplex twiddle = make_cuComplex(realPart, imagPart);
            temp = cuCaddf(temp, cuCmulf(data[2*N + k], twiddle));
        }
        data[idx] = temp;
    }
}

// 三模冗余投票函数
cuComplex tmr_vote(cuComplex a, cuComplex b, cuComplex c) {
    // 比较实部
    bool ab_real = (fabs(cuCrealf(a) - cuCrealf(b)) < 1e-12);
    bool ac_real = (fabs(cuCrealf(a) - cuCrealf(c)) < 1e-12);
    bool bc_real = (fabs(cuCrealf(b) - cuCrealf(c)) < 1e-12);
    
    // 比较虚部
    bool ab_imag = (fabs(cuCimagf(a) - cuCimagf(b)) < 1e-12);
    bool ac_imag = (fabs(cuCimagf(a) - cuCimagf(c)) < 1e-12);
    bool bc_imag = (fabs(cuCimagf(b) - cuCimagf(c)) < 1e-12);

    // 实部表决
    float real_result;
    if (ab_real && ac_real) {
        real_result = cuCrealf(a);
    } else if (ab_real) {
        real_result = cuCrealf(a);
    } else if (ac_real) {
        real_result = cuCrealf(a);
    } else if (bc_real) {
        real_result = cuCrealf(b);
    } else {
        // 三个都不一致，取中值
        float vals[3] = {cuCrealf(a), cuCrealf(b), cuCrealf(c)};
        if (vals[0] > vals[1]) std::swap(vals[0], vals[1]);
        if (vals[1] > vals[2]) std::swap(vals[1], vals[2]);
        if (vals[0] > vals[1]) std::swap(vals[0], vals[1]);
        real_result = vals[1];
    }

    // 虚部表决
    float imag_result;
    if (ab_imag && ac_imag) {
        imag_result = cuCimagf(a);
    } else if (ab_imag) {
        imag_result = cuCimagf(a);
    } else if (ac_imag) {
        imag_result = cuCimagf(a);
    } else if (bc_imag) {
        imag_result = cuCimagf(b);
    } else {
        // 三个都不一致，取中值
        float vals[3] = {cuCimagf(a), cuCimagf(b), cuCimagf(c)};
        if (vals[0] > vals[1]) std::swap(vals[0], vals[1]);
        if (vals[1] > vals[2]) std::swap(vals[1], vals[2]);
        if (vals[0] > vals[1]) std::swap(vals[0], vals[1]);
        imag_result = vals[1];
    }

    return make_cuComplex(real_result, imag_result);
}

int main() {
    const int N = 8;
    
    // 分配3N大小的数组
    cuComplex* h_data = new cuComplex[TOTAL_MODULES * N];
    cuComplex* h_final_result = new cuComplex[N];

    // 初始化三组相同数据
    for (int i = 0; i < N; i++) {
        h_data[i] = make_cuComplex(1.0f, 0.0f);          // 第一组
        h_data[N + i] = make_cuComplex(1.0f, 0.0f);       // 第二组
        h_data[2*N + i] = make_cuComplex(1.0f, 0.0f);     // 第三组
    }

    cuComplex* d_data;
    cudaMalloc((void**)&d_data, TOTAL_MODULES * N * sizeof(cuComplex));
    cudaMemcpy(d_data, h_data, TOTAL_MODULES * N * sizeof(cuComplex), cudaMemcpyHostToDevice);

    // 启动FFT kernel，使用3个block处理三组数据
    fft_kernel<<<3, N>>>(d_data, N);

    // 拷贝结果回主机
    cudaMemcpy(h_data, d_data, TOTAL_MODULES * N * sizeof(cuComplex), cudaMemcpyDeviceToHost);

    // 执行三模冗余表决
    for (int i = 0; i < N; i++) {
        h_final_result[i] = tmr_vote(h_data[i], h_data[N + i], h_data[2*N + i]);
    }

    // 错误检测和结果输出
    int error_count = 0;
    for (int i = 0; i < N; i++) {
        bool real_match = (fabs(cuCrealf(h_data[i]) - cuCrealf(h_data[N + i])) < 1e-12) &&
                          (fabs(cuCrealf(h_data[i]) - cuCrealf(h_data[2*N + i])) < 1e-12);
        bool imag_match = (fabs(cuCimagf(h_data[i]) - cuCimagf(h_data[N + i])) < 1e-12) &&
                          (fabs(cuCimagf(h_data[i]) - cuCimagf(h_data[2*N + i])) < 1e-12);
        
        if (!real_match || !imag_match) {
            error_count++;
            std::cout << "Error detected at index " << i << ":\n"
                      << "  Module1: " << cuCrealf(h_data[i]) << "+" << cuCimagf(h_data[i]) << "i\n"
                      << "  Module2: " << cuCrealf(h_data[N + i]) << "+" << cuCimagf(h_data[N + i]) << "i\n"
                      << "  Module3: " << cuCrealf(h_data[2*N + i]) << "+" << cuCimagf(h_data[2*N + i]) << "i\n"
                      << "  Voted: " << cuCrealf(h_final_result[i]) << "+" << cuCimagf(h_final_result[i]) << "i\n";
        }
    }

    if (error_count == 0) {
        std::cout << "All three modules produced identical results:\n";
    } else {
        std::cout << "\nTotal errors detected and corrected: " << error_count << "\n";
    }

    // 输出最终结果
    for (int i = 0; i < N; i++) {
        std::cout << "Final[" << i << "] = " << cuCrealf(h_final_result[i]) 
                  << " + " << cuCimagf(h_final_result[i]) << "i\n";
    }

    // 清理资源
    delete[] h_data;
    delete[] h_final_result;
    cudaFree(d_data);

    return 0;
}